
import sqlite3

from util import dict_factory

class DataBaseFunction:

    def __init__(self) -> None:
        
        self.conn = sqlite3.connect('sqlite3.db')
        self.conn.row_factory = dict_factory
        # 创建一个游标对象
        self.cursor = self.conn.cursor()

    def __enter__(self):

        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):

        if exc_type is None:
            self.conn.commit()  # 如果没有异常发生，提交事务
        self.cursor.close()
        self.conn.close()

    def selete_tasks_by_status(self, status="success"): 
        '''输入状态查找任务'''

        sql = "select * from tasks where status=?"

        self.cursor.execute(sql, (status, ))

        return self.cursor.fetchall()
    
    def selete_tasks_by_status_complete(self): 
        '''找到完成的任务：success, fail'''

        sql = "select * from tasks where status='success' or status='fail'"

        self.cursor.execute(sql)

        return self.cursor.fetchall()
    
    def update_tasks_status(self, task_id, status):
        '''升级tasks表状态'''

        sql = "update tasks set status=? where task_id=?"

        self.cursor.execute(sql, (status, task_id))

    def update_tasks_return_data(self, task_id, return_data):
        '''更新指定task_id的return_data字段'''

        sql = "update tasks set return_data=? where task_id=?"

        self.cursor.execute(sql, (return_data, task_id))

    def selete_downlink_tasks_by_tasks_table_id(self, tasks_table_id): 
        '''通过tasks_id查找downlink表中的任务'''

        sql = "select * from downlink_tasks where task_id=?"

        self.cursor.execute(sql, (tasks_table_id, ))

        return self.cursor.fetchall()

    def update_downlink_tasks_status(self, task_id, status):
        '''升级downlink_tasks表状态'''

        sql = "update downlink_tasks set status=? where id=?"

        self.cursor.execute(sql, (status, task_id))

    def select_downlink_task_id_by_status_execute_and_sensor_mac(self, sensor_mac, sn):
        '''查找senor_mac对应的sn号其当前状态为execute的任务id'''

        sql = "select id from downlink_tasks where status='execute' and sensor_mac=? and sn=?"

        self.cursor.execute(sql, (sensor_mac, sn))

        result = self.cursor.fetchone()

        if result:

            return result["id"]
        else:

            return None

    def update_downlink_return_data(self, return_data, task_id):
        '''升级dowlink_tasks 表 的return data 和状态'''

        sql = "update downlink_tasks set return_data=? , status='success' where id=?"

        self.cursor.execute(sql, (return_data, task_id))

    def get_sensor_update_task_list(self, base_station_mac)->list:
        '''获取传感器升级任务列表'''

        sql = "select count(*) from  downlink_tasks where base_station_mac=? and status='execute'"

        self.cursor.execute(sql, (base_station_mac, ))

        quantity_in_execution = self.cursor.fetchone()['count']

        if quantity_in_execution < 6: # 同一基站最多同时执行6个传感器升级任务

            sql = "select * from downlink_task where status='wait' and base_station_mac=? limit ?"

            self.cursor.execute(sql, (base_station_mac, 6-quantity_in_execution))

            return self.cursor.fetchall()
        
        else:

            return []
        
    def select_base_station_mac_by_sensor_mac(self, sensor_mac):
        '''传感的mac找到对应基站的mac'''

        sql = """\
        SELECT bs.mac AS base_station_mac
            FROM base_stations bs
            JOIN brackets b ON bs.id = b.base_station_id
            JOIN sensors s ON b.id = s.brackets_id
            WHERE s.mac = ?;
        """

        self.cursor.execute(sql, (sensor_mac, ))

        result = self.cursor.fetchone()

        if result:

            return result["base_station_mac"]
        else:

            return None
        
    def select_tasks_by_task_type(self, task_type, status='wait'):
        '''以任务类型和状态查询主任务'''

        sql = "select * from tasks where task_type=? and status=?"

        self.cursor.execute(sql, (task_type, status))

        return self.cursor.fetchall()
    
    def select_base_station_mac_by_sensor_upgrade(self):
        '''查找当前有传感器升级任务的base_station_mac'''
        sql = """\
            SELECT DISTINCT base_station_mac FROM downlink_tasks WHERE task_id IN (
                SELECT task_id FROM tasks WHERE task_type = 'sensor_upgrade' AND status = 'wait'
        );\
        """

        self.cursor.execute(sql)

        return self.cursor.fetchall()
    
if __name__ == "__main__":

    with DataBaseFunction() as db:

        db.update_tasks_status(33, 'success')

        v = db.selete_success_tasks()

    print(v)